#include "common.h"
#include "kernel_operator.h"
#include <cstdint>
#ifdef __CCE_KT_TEST__
#include <iostream>
#define DTYPE_X half;
#define DTYPE_Y half;
#endif
namespace AscendC
{
template <typename _DT_X = half, typename _DT_Y = half> class KernelOtherHalf
{
  public:
    TPipe pipe;
    int32_t last_size = 0;
    float cof_p = 0;
    float cof_esp = 0;
    CommonTiling tilingData;
    NdTensorSortByDimHelper *sortHelper = nullptr;
    int64_t big_loop_count = 1;
    DefInTensor(X);
    DefOutTensor(Y);

    DefBufVECCALC(TEMP);
    DefBufVECCALC(CAST_X);
    DefBufVECCALC(SUM_VALUE);
    DefBufVECCALC(WORK_LOCAL);

  public:
    __aicore__ inline KernelOtherHalf()
    {
    }
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, int64_t size, float p, float epsilon, CommonTiling &tiling,
                                NdTensorSortByDimHelper *sort_helper)
    {
        tilingData = tiling;
        sortHelper = sort_helper;
        if (sortHelper)
        {
            big_loop_count = tilingData.blockLength;
            this->last_size = sortHelper->sortByDimDimSize % tilingData.bufferSize;
            if (this->last_size == 0)
            {
                this->last_size = tilingData.bufferSize;
            }
            tilingData.loopCount = CEIL_DIV(sortHelper->sortByDimDimSize, tilingData.bufferSize);

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx,
                            ALIGN_TO((tilingData.blockLength * sortHelper->sortByDimDimSize), (32 / sizeof(_DT_X))));
            SetGlobalBuffer(Y, y, tilingData.startIdx, ALIGN_TO(tilingData.blockLength, (32 / sizeof(_DT_X))));
        }
        else
        {
            big_loop_count = 1;
            this->last_size = tilingData.last_size;

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx, tilingData.blockLength);
            SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        }
        cof_p = p;
        cof_esp = epsilon;
        // 初始化队列
        InitQueueSimple(X, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP, tilingData.bufferSize * sizeof(float));
        InitTBufBuffer(CAST_X, tilingData.bufferSize * sizeof(float));
        InitTBufBuffer(SUM_VALUE, tilingData.bufferSize * sizeof(float));
        InitTBufBuffer(WORK_LOCAL, tilingData.bufferSize * sizeof(float));
    }
    __aicore__ inline void Process()
    {
        for (int64_t j = 0; j < big_loop_count; ++j)
        {
            // SUM_VALUE = 0
            TBufGet(SUM_VALUE, _DT_X);
            TBufTensorName(SUM_VALUE).SetValue(0, _DT_X(0));

            auto finnal_progress = tilingData.loopCount - 1;
            for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
            {
                CopyIn(i, tilingData.bufferSize, false);
                Compute(i, tilingData.bufferSize, false);
                CopyOut(i, tilingData.bufferSize, false, j);
            }
            auto copy_count = ALIGN_TO(this->last_size, 32 / sizeof(TypeOf(X)));
            CopyIn(finnal_progress, copy_count, true);
            Compute(finnal_progress, tilingData.finalCalcCount, true);
            CopyOut(finnal_progress, copy_count, true, j);
        }
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        if (sortHelper)
        {
            QueAllocSimple(X);
            int32_t count = is_final ? this->last_size : calcCount;
            sortHelper->read_to_tensor(LTensorName(X), count);
            EnQue(X);
        }
        else
        {
            auto index = progress * tilingData.bufferSize;
            EnQueGlobal2Local(X, index, calcCount);
        }
    }
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        DeQueSimple(X);
        TBufGet(TEMP, float);
        TBufGet(CAST_X, float);
        TBufGet(SUM_VALUE, float);
        TBufGet(WORK_LOCAL, float);
        Cast(TBufTensorName(CAST_X), LTensorName(X), RoundMode::CAST_NONE, calcCount);
        // abs(x)
        Abs(TBufTensorName(TEMP), TBufTensorName(CAST_X), calcCount);
        // abs(x) ^ p
        MyPow(TBufTensorName(CAST_X), TBufTensorName(TEMP), float(cof_p), calcCount);
        // sum (abs(x) ^ p)
        int32_t count = is_final ? this->last_size : calcCount;
        ReduceSum(TBufTensorName(TEMP), TBufTensorName(CAST_X), TBufTensorName(WORK_LOCAL), count);
        // 加到总数上
        Add(TBufTensorName(SUM_VALUE), TBufTensorName(TEMP), TBufTensorName(SUM_VALUE), 1);
        // sum ^ (1/p)
        if (is_final)
        {
            QueAllocSimple(Y);
            MyPow(TBufTensorName(CAST_X), TBufTensorName(SUM_VALUE), float(1 / (cof_p + cof_esp)), 1);
            Cast(LTensorName(Y), TBufTensorName(CAST_X), RoundMode::CAST_NONE, 1);
            EnQue(Y);
        }
        QueFree(X);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount, bool is_final, int64_t big_index)
    {
        if (is_final)
        {
            DeQueLocal2Global(Y, big_index, 32 / sizeof(_DT_X));
        }
    }
};

template <typename _DT_X = float, typename _DT_Y = float> class KernelOtherFloat
{
  public:
    TPipe pipe;
    int32_t last_size = 0;
    float cof_p = 0;
    float cof_esp = 0;
    CommonTiling tilingData;
    NdTensorSortByDimHelper *sortHelper = nullptr;
    int64_t big_loop_count = 1;
    DefInTensor(X);
    DefOutTensor(Y);

    DefBufVECCALC(TEMP);
    DefBufVECCALC(SUM_VALUE);
    DefBufVECCALC(WORK_LOCAL);

  public:
    __aicore__ inline KernelOtherFloat()
    {
    }
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, int64_t size, float p, float epsilon, CommonTiling &tiling,
                                NdTensorSortByDimHelper *sort_helper = nullptr)
    {
        tilingData = tiling;
        cof_p = p;
        cof_esp = epsilon;
        sortHelper = sort_helper;
        if (sortHelper)
        {
            big_loop_count = tilingData.blockLength;
            this->last_size = sortHelper->sortByDimDimSize % tilingData.bufferSize;
            if (this->last_size == 0)
            {
                this->last_size = tilingData.bufferSize;
            }
            tilingData.loopCount = CEIL_DIV(sortHelper->sortByDimDimSize, tilingData.bufferSize);

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx,
                            ALIGN_TO((tilingData.blockLength * sortHelper->sortByDimDimSize), (32 / sizeof(_DT_X))));
            SetGlobalBuffer(Y, y, tilingData.startIdx, ALIGN_TO(tilingData.blockLength, (32 / sizeof(_DT_X))));
        }
        else
        {
            big_loop_count = 1;
            this->last_size = tilingData.last_size;
            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx, tilingData.blockLength);
            SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        }
        // 初始化队列
        InitQueueSimple(X, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(SUM_VALUE, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(WORK_LOCAL, tilingData.bufferSize * sizeof(_DT_X));
    }
    __aicore__ inline void Process()
    {
        for (int64_t j = 0; j < big_loop_count; ++j)
        {
            // SUM_VALUE = 0
            TBufGet(SUM_VALUE, _DT_X);
            TBufTensorName(SUM_VALUE).SetValue(0, _DT_X(0));

            auto finnal_progress = tilingData.loopCount - 1;
            for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
            {
                CopyIn(i, tilingData.bufferSize, false);
                Compute(i, tilingData.bufferSize, false);
                CopyOut(i, tilingData.bufferSize, false, j);
            }
            CopyIn(finnal_progress, tilingData.finalCalcCount, true);
            Compute(finnal_progress, tilingData.finalCalcCount, true);
            CopyOut(finnal_progress, tilingData.finalCalcCount, true, j);
        }
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        if (sortHelper)
        {
            QueAllocSimple(X);
            int32_t count = is_final ? this->last_size : calcCount;
            sortHelper->read_to_tensor(LTensorName(X), count);
            EnQue(X);
        }
        else
        {
            auto index = progress * tilingData.bufferSize;
            EnQueGlobal2Local(X, index, calcCount);
        }
    }
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        DeQueSimple(X);
        TBufGet(TEMP, _DT_X);
        TBufGet(SUM_VALUE, _DT_X);
        TBufGet(WORK_LOCAL, _DT_X);
        // abs(x)
        Abs(TBufTensorName(TEMP), LTensorName(X), calcCount);
        // abs(x) ^ p
        MyPow(LTensorName(X), TBufTensorName(TEMP), _DT_X(cof_p), calcCount);
        // sum (abs(x) ^ p)
        int32_t count = is_final ? this->last_size : calcCount;
        ReduceSum(TBufTensorName(TEMP), LTensorName(X), TBufTensorName(WORK_LOCAL), count);
        // 加到总数上
        Add(TBufTensorName(SUM_VALUE), TBufTensorName(SUM_VALUE), TBufTensorName(TEMP), 1);
        // sum ^ (1/p)
        if (is_final)
        {
            QueAllocSimple(Y);
            MyPow(LTensorName(Y), TBufTensorName(SUM_VALUE), _DT_X(1 / (cof_p + cof_esp)), 1);
            EnQue(Y);
        }
        QueFree(X);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount, bool is_final, int64_t big_index)
    {
        if (is_final)
        {
            DeQueLocal2Global(Y, big_index, 32 / sizeof(_DT_X));
        }
    }
};
template <typename _DT_X, typename _DT_Y> class KernelP2
{
  public:
    TPipe pipe;
    int32_t last_size = 0;
    CommonTiling tilingData;
    NdTensorSortByDimHelper *sort_helper = nullptr;
    int64_t big_loop_count = 1;
    DefInTensor(X);
    DefOutTensor(Y);

    DefBufVECCALC(TEMP);
    DefBufVECCALC(SUM_VALUE);
    DefBufVECCALC(WORK_LOCAL);

  public:
    __aicore__ inline KernelP2()
    {
    }
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, int64_t size, CommonTiling &tiling,
                                NdTensorSortByDimHelper *sortHelper = nullptr)
    {
        tilingData = tiling;
        sort_helper = sortHelper;

        if (sortHelper)
        {
            big_loop_count = tilingData.blockLength;
            this->last_size = sortHelper->sortByDimDimSize % tilingData.bufferSize;
            if (this->last_size == 0)
            {
                this->last_size = tilingData.bufferSize;
            }
            tilingData.loopCount = CEIL_DIV(sortHelper->sortByDimDimSize, tilingData.bufferSize);

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx,
                            ALIGN_TO((tilingData.blockLength * sortHelper->sortByDimDimSize), (32 / sizeof(_DT_X))));
            SetGlobalBuffer(Y, y, tilingData.startIdx, ALIGN_TO(tilingData.blockLength, (32 / sizeof(_DT_X))));
        }
        else
        {
            big_loop_count = 1;
            this->last_size = tilingData.last_size;

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx, tilingData.blockLength);
            SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        }
        // 初始化队列
        InitQueueSimple(X, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(SUM_VALUE, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(WORK_LOCAL, tilingData.bufferSize * sizeof(_DT_X));
    }
    __aicore__ inline void Process()
    {
        for (int64_t j = 0; j < big_loop_count; ++j)
        {
            // SUM_VALUE = 0
            TBufGet(SUM_VALUE, _DT_X);
            TBufTensorName(SUM_VALUE).SetValue(0, _DT_X(0));

            auto finnal_progress = tilingData.loopCount - 1;
            for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
            {
                CopyIn(i, tilingData.bufferSize, false);
                Compute(i, tilingData.bufferSize, false);
                CopyOut(i, tilingData.bufferSize, false, j);
            }
            CopyIn(finnal_progress, tilingData.finalCalcCount, true);
            Compute(finnal_progress, tilingData.finalCalcCount, true);
            CopyOut(finnal_progress, tilingData.finalCalcCount, true, j);
        }
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount, bool is_final)
    {

        if (sort_helper)
        {
            QueAllocSimple(X);
            int32_t count = is_final ? this->last_size : calcCount;
            sort_helper->read_to_tensor(LTensorName(X), count);
            EnQue(X);
        }
        else
        {
            auto index = progress * tilingData.bufferSize;
            EnQueGlobal2Local(X, index, calcCount);
        }
    }
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        DeQueSimple(X);
        TBufGet(TEMP, _DT_X);
        TBufGet(SUM_VALUE, _DT_X);
        TBufGet(WORK_LOCAL, _DT_X);
        // abs(x)
        Abs(TBufTensorName(TEMP), LTensorName(X), calcCount);
        // abs(x) ^ 2
        Mul(LTensorName(X), TBufTensorName(TEMP), TBufTensorName(TEMP), calcCount);
        // sum (abs(x) ^ 2)
        int32_t count = is_final ? this->last_size : calcCount;
        ReduceSum(TBufTensorName(TEMP), LTensorName(X), TBufTensorName(WORK_LOCAL), count);
        // 加到总数上
        Add(TBufTensorName(SUM_VALUE), TBufTensorName(SUM_VALUE), TBufTensorName(TEMP), 1);
        // sum ^ (1/p)
        if (is_final)
        {
            QueAllocSimple(Y);
            Sqrt(LTensorName(Y), TBufTensorName(SUM_VALUE), 1);
            EnQue(Y);
        }
        QueFree(X);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount, bool is_final, int64_t big_index)
    {
        if (is_final)
        {
            DeQueLocal2Global(Y, big_index, 32 / sizeof(_DT_X));
        }
    }
};

template <typename _DT_X, typename _DT_Y> class KernelP1
{
  public:
    TPipe pipe;
    int32_t last_size = 0;
    CommonTiling tilingData;
    NdTensorSortByDimHelper *sort_helper = nullptr;
    int64_t big_loop_count = 1;
    DefInTensor(X);
    DefOutTensor(Y);

    DefBufVECCALC(TEMP);
    DefBufVECCALC(SUM_VALUE);
    DefBufVECCALC(WORK_LOCAL);

  public:
    __aicore__ inline KernelP1()
    {
    }
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, int64_t size, CommonTiling &tiling,
                                NdTensorSortByDimHelper *sortHelper = nullptr)
    {
        tilingData = tiling;
        sort_helper = sortHelper;

        if (sortHelper)
        {
            big_loop_count = tilingData.blockLength;
            this->last_size = sortHelper->sortByDimDimSize % tilingData.bufferSize;
            if (this->last_size == 0)
            {
                this->last_size = tilingData.bufferSize;
            }
            tilingData.loopCount = CEIL_DIV(sortHelper->sortByDimDimSize, tilingData.bufferSize);

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx,
                            ALIGN_TO((tilingData.blockLength * sortHelper->sortByDimDimSize), (32 / sizeof(_DT_X))));
            SetGlobalBuffer(Y, y, tilingData.startIdx, ALIGN_TO(tilingData.blockLength, (32 / sizeof(_DT_X))));
        }
        else
        {
            big_loop_count = 1;
            this->last_size = tilingData.last_size;

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx, tilingData.blockLength);
            SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        }

        // 初始化队列
        InitQueueSimple(X, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(SUM_VALUE, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(WORK_LOCAL, tilingData.bufferSize * sizeof(_DT_X));
    }
    __aicore__ inline void Process()
    {
        for (int64_t j = 0; j < big_loop_count; ++j)
        {
            // SUM_VALUE = 0
            TBufGet(SUM_VALUE, _DT_X);
            TBufTensorName(SUM_VALUE).SetValue(0, _DT_X(0));

            int64_t finnal_progress = tilingData.loopCount - 1;
            for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
            {
                CopyIn(i, tilingData.bufferSize, false);
                Compute(i, tilingData.bufferSize, false);
                CopyOut(i, tilingData.bufferSize, false, j);
            }
            CopyIn(finnal_progress, tilingData.finalCalcCount, true);
            Compute(finnal_progress, tilingData.finalCalcCount, true);
            CopyOut(finnal_progress, tilingData.finalCalcCount, true, j);
        }
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        if (sort_helper)
        {
            QueAllocSimple(X);
            int32_t count = is_final ? this->last_size : calcCount;
            sort_helper->read_to_tensor(LTensorName(X), count);
            EnQue(X);
        }
        else
        {
            auto index = progress * tilingData.bufferSize;
            EnQueGlobal2Local(X, index, calcCount);
        }
    }
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        DeQueSimple(X);
        TBufGet(TEMP, _DT_X);
        TBufGet(SUM_VALUE, _DT_X);
        TBufGet(WORK_LOCAL, _DT_X);
        PRINT_TENSOR(LTensorName(X), 8, progress);
        // abs(x)
        Abs(TBufTensorName(TEMP), LTensorName(X), calcCount);
        // sum (abs(x))
        int32_t count = is_final ? this->last_size : calcCount;
        ReduceSum(LTensorName(X), TBufTensorName(TEMP), TBufTensorName(WORK_LOCAL), count);
        // 加到总数上
        Add(TBufTensorName(SUM_VALUE), TBufTensorName(SUM_VALUE), LTensorName(X), 1);
        if (is_final)
        {
            QueAllocSimple(Y);
            DataCopy(LTensorName(Y), LTensorName(X), 32 / sizeof(_DT_X));
            EnQue(Y);
        }
        QueFree(X);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount, bool is_final, int64_t big_index)
    {
        if (is_final)
        {
            DeQueLocal2Global(Y, big_index, 32 / sizeof(_DT_X));
        }
    }
};

template <typename _DT_X, typename _DT_Y> class KernelP0
{
  public:
    TPipe pipe;
    int32_t last_size = 0;
    CommonTiling tilingData;
    NdTensorSortByDimHelper *sort_helper = nullptr;
    int64_t big_loop_count = 1;
    DefInTensor(X);
    DefOutTensor(Y);

    DefBufVECCALC(TEMP);
    DefBufVECCALC(TEMP_0);
    DefBufVECCALC(TEMP_1);
    DefBufVECCALC(SUM_VALUE);
    DefBufVECCALC(WORK_LOCAL);

  public:
    __aicore__ inline KernelP0()
    {
    }
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, int64_t size, CommonTiling &tiling,
                                NdTensorSortByDimHelper *sortHelper = nullptr)
    {
        tilingData = tiling;
        sort_helper = sortHelper;

        if (sortHelper)
        {
            big_loop_count = tilingData.blockLength;
            this->last_size = sortHelper->sortByDimDimSize % tilingData.bufferSize;
            if (this->last_size == 0)
            {
                this->last_size = tilingData.bufferSize;
            }
            tilingData.loopCount = CEIL_DIV(sortHelper->sortByDimDimSize, tilingData.bufferSize);

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx,
                            ALIGN_TO((tilingData.blockLength * sortHelper->sortByDimDimSize), (32 / sizeof(_DT_X))));
            SetGlobalBuffer(Y, y, tilingData.startIdx, ALIGN_TO(tilingData.blockLength, (32 / sizeof(_DT_X))));
        }
        else
        {
            big_loop_count = 1;
            this->last_size = tilingData.last_size;

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx, tilingData.blockLength);
            SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        }
        // 初始化队列
        InitQueueSimple(X, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(TEMP_0, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(TEMP_1, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(SUM_VALUE, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(WORK_LOCAL, tilingData.bufferSize * sizeof(_DT_X));

        // TEMP_0 = 0
        TBufGet(TEMP_0, _DT_X);
        for (auto i = 0; i < tilingData.bufferSize; ++i)
        {
            TBufTensorName(TEMP_0).SetValue(i, _DT_X(0));
        }
        // TEMP_1 = 1
        TBufGet(TEMP_1, _DT_X);
        for (auto i = 0; i < tilingData.bufferSize; ++i)
        {
            TBufTensorName(TEMP_1).SetValue(i, _DT_X(1));
        }
    }
    __aicore__ inline void Process()
    {
        for (int64_t j = 0; j < big_loop_count; ++j)
        {
            // SUM_VALUE = 0
            TBufGet(SUM_VALUE, _DT_X);
            TBufTensorName(SUM_VALUE).SetValue(0, _DT_X(0));

            int64_t finnal_progress = tilingData.loopCount - 1;
            for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
            {
                CopyIn(i, tilingData.bufferSize, false);
                Compute(i, tilingData.bufferSize, false);
                CopyOut(i, tilingData.bufferSize, false, j);
            }
            CopyIn(finnal_progress, tilingData.finalCalcCount, true);
            Compute(finnal_progress, tilingData.finalCalcCount, true);
            CopyOut(finnal_progress, tilingData.finalCalcCount, true, j);
        }
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        if (sort_helper)
        {
            QueAllocSimple(X);
            int32_t count = is_final ? this->last_size : calcCount;
            sort_helper->read_to_tensor(LTensorName(X), count);
            EnQue(X);
        }
        else
        {
            auto index = progress * tilingData.bufferSize;
            EnQueGlobal2Local(X, index, calcCount);
        }
    }
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        DeQueSimple(X);
        TBufGet(TEMP, int8_t);
        TBufGet(TEMP_0, _DT_X);
        TBufGet(TEMP_1, _DT_X);
        TBufGet(SUM_VALUE, _DT_X);
        TBufGet(WORK_LOCAL, _DT_X);
        // x != 0
        Compare(TBufTensorName(TEMP), LTensorName(X), TBufTensorName(TEMP_0), CMPMODE::NE, calcCount);
        // x = x != 0 ? 1 : 0
        Select(LTensorName(X), TBufTensorName(TEMP), TBufTensorName(TEMP_1), _DT_X(0), SELMODE::VSEL_TENSOR_SCALAR_MODE,
               calcCount);
        // sum (x != 0)
        auto type_temp = TBufTensorName(TEMP).ReinterpretCast<_DT_X>();
        int32_t count = is_final ? this->last_size : calcCount;
        ReduceSum(type_temp, LTensorName(X), TBufTensorName(WORK_LOCAL), count);
        // 加到总数上
        Add(TBufTensorName(SUM_VALUE), TBufTensorName(SUM_VALUE), type_temp, 1);
        // sum ^ (1/p)
        if (is_final)
        {
            QueAllocSimple(Y);
            DataCopy(LTensorName(Y), TBufTensorName(SUM_VALUE), 32 / sizeof(_DT_X));
            EnQue(Y);
        }
        QueFree(X);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount, bool is_final, int64_t big_index)
    {
        if (is_final)
        {
            DeQueLocal2Global(Y, big_index, 32 / sizeof(_DT_X));
        }
    }
};

template <typename _DT_X, typename _DT_Y> class KernelPInf
{
  public:
    TPipe pipe;
    int32_t last_size = 0;
    bool is_first_compare = true;
    CommonTiling tilingData;
    NdTensorSortByDimHelper *sort_helper = nullptr;
    int64_t big_loop_count = 1;
    DefInTensor(X);
    DefOutTensor(Y);

    DefBufVECCALC(TEMP);
    DefBufVECCALC(STORE_VALUE);
    DefBufVECCALC(WORK_LOCAL);

  public:
    __aicore__ inline KernelPInf()
    {
    }
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, int64_t size, CommonTiling &tiling,
                                NdTensorSortByDimHelper *sortHelper)
    {
        tilingData = tiling;
        sort_helper = sortHelper;

        if (sortHelper)
        {
            big_loop_count = tilingData.blockLength;
            this->last_size = sortHelper->sortByDimDimSize % tilingData.bufferSize;
            if (this->last_size == 0)
            {
                this->last_size = tilingData.bufferSize;
            }
            tilingData.loopCount = CEIL_DIV(sortHelper->sortByDimDimSize, tilingData.bufferSize);

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx,
                            ALIGN_TO((tilingData.blockLength * sortHelper->sortByDimDimSize), (32 / sizeof(_DT_X))));
            SetGlobalBuffer(Y, y, tilingData.startIdx, ALIGN_TO(tilingData.blockLength, (32 / sizeof(_DT_X))));
        }
        else
        {
            big_loop_count = 1;
            this->last_size = tilingData.last_size;

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx, tilingData.blockLength);
            SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        }
        // 初始化队列
        InitQueueSimple(X, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(STORE_VALUE, (tilingData.bufferSize + 1) * sizeof(_DT_X));
        InitTBufBuffer(WORK_LOCAL, tilingData.bufferSize * sizeof(_DT_X));
    }
    __aicore__ inline void Process()
    {
        for (int64_t j = 0; j < big_loop_count; ++j)
        {
            this->is_first_compare = true;
            int64_t finnal_progress = tilingData.loopCount - 1;
            for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
            {
                CopyIn(i, tilingData.bufferSize, false);
                Compute(i, tilingData.bufferSize, false);
                CopyOut(i, tilingData.bufferSize, false, j);
            }
            CopyIn(finnal_progress, tilingData.finalCalcCount, true);
            Compute(finnal_progress, tilingData.finalCalcCount, true);
            CopyOut(finnal_progress, tilingData.finalCalcCount, true, j);
        }
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        if (sort_helper)
        {
            QueAllocSimple(X);
            int32_t count = is_final ? this->last_size : calcCount;
            sort_helper->read_to_tensor(LTensorName(X), count);
            EnQue(X);
        }
        else
        {
            auto index = progress * tilingData.bufferSize;
            EnQueGlobal2Local(X, index, calcCount);
        }
    }
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        DeQueSimple(X);
        TBufGet(TEMP, _DT_X);
        TBufGet(STORE_VALUE, _DT_X);
        TBufGet(WORK_LOCAL, _DT_X);
        // abs(x)
        Abs(TBufTensorName(TEMP), LTensorName(X), calcCount);
        // reduce max
        int32_t count = is_final ? this->last_size : calcCount;
        if (is_first_compare)
        {
            ReduceMax(TBufTensorName(STORE_VALUE), TBufTensorName(TEMP), TBufTensorName(WORK_LOCAL), count);
            is_first_compare = false;
        }
        else
        {
            DataCopy(TBufTensorName(STORE_VALUE)[1], TBufTensorName(TEMP), calcCount);
            ReduceMax(TBufTensorName(STORE_VALUE), TBufTensorName(STORE_VALUE), TBufTensorName(WORK_LOCAL), count + 1);
        }
        if (is_final)
        {
            QueAllocSimple(Y);
            DataCopy(LTensorName(Y), TBufTensorName(STORE_VALUE), 32 / sizeof(_DT_X));
            EnQue(Y);
        }
        QueFree(X);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount, bool is_final, int64_t big_index)
    {
        if (is_final)
        {
            DeQueLocal2Global(Y, big_index, 32 / sizeof(_DT_X));
        }
    }
};

template <typename _DT_X, typename _DT_Y> class KernelPnInf
{
  public:
    TPipe pipe;
    int32_t last_size = 0;
    bool is_first_compare = true;
    CommonTiling tilingData;
    NdTensorSortByDimHelper *sort_helper = nullptr;
    int64_t big_loop_count = 1;
    DefInTensor(X);
    DefOutTensor(Y);

    DefBufVECCALC(TEMP);
    DefBufVECCALC(STORE_VALUE);
    DefBufVECCALC(WORK_LOCAL);

  public:
    __aicore__ inline KernelPnInf()
    {
    }
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, int64_t size, CommonTiling &tiling,
                                NdTensorSortByDimHelper *sortHelper = nullptr)
    {
        tilingData = tiling;
        sort_helper = sortHelper;

        if (sort_helper)
        {
            big_loop_count = tilingData.blockLength;
            this->last_size = sortHelper->sortByDimDimSize % tilingData.bufferSize;
            if (this->last_size == 0)
            {
                this->last_size = tilingData.bufferSize;
            }
            tilingData.loopCount = CEIL_DIV(sortHelper->sortByDimDimSize, tilingData.bufferSize);

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx,
                            ALIGN_TO((tilingData.blockLength * sortHelper->sortByDimDimSize), (32 / sizeof(_DT_X))));
            SetGlobalBuffer(Y, y, tilingData.startIdx, ALIGN_TO(tilingData.blockLength, (32 / sizeof(_DT_X))));
        }
        else
        {
            big_loop_count = 1;
            this->last_size = tilingData.last_size;

            // 切分global
            SetGlobalBuffer(X, x, tilingData.startIdx, tilingData.blockLength);
            SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        }
        // 初始化队列
        InitQueueSimple(X, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(STORE_VALUE, (tilingData.bufferSize + 1) * sizeof(_DT_X));
        InitTBufBuffer(WORK_LOCAL, tilingData.bufferSize * sizeof(_DT_X));
    }
    __aicore__ inline void Process()
    {
        for (int64_t j = 0; j < big_loop_count; ++j)
        {
            this->is_first_compare = true;
            int64_t finnal_progress = tilingData.loopCount - 1;
            for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
            {
                CopyIn(i, tilingData.bufferSize, false);
                Compute(i, tilingData.bufferSize, false);
                CopyOut(i, tilingData.bufferSize, false, j);
            }
            CopyIn(finnal_progress, tilingData.finalCalcCount, true);
            Compute(finnal_progress, tilingData.finalCalcCount, true);
            CopyOut(finnal_progress, tilingData.finalCalcCount, true, j);
        }
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        if (sort_helper)
        {
            QueAllocSimple(X);
            int32_t count = is_final ? this->last_size : calcCount;
            sort_helper->read_to_tensor(LTensorName(X), count);
            EnQue(X);
        }
        else
        {
            auto index = progress * tilingData.bufferSize;
            EnQueGlobal2Local(X, index, calcCount);
        }
    }
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        DeQueSimple(X);
        TBufGet(TEMP, _DT_X);
        TBufGet(STORE_VALUE, _DT_X);
        TBufGet(WORK_LOCAL, _DT_X);
        // abs(x)
        Abs(TBufTensorName(TEMP), LTensorName(X), calcCount);
        // reduce max
        int32_t count = is_final ? this->last_size : calcCount;
        if (is_first_compare)
        {
            ReduceMin(TBufTensorName(STORE_VALUE), TBufTensorName(TEMP), TBufTensorName(WORK_LOCAL), count);
            is_first_compare = false;
        }
        else
        {
            DataCopy(TBufTensorName(STORE_VALUE)[1], TBufTensorName(TEMP), calcCount);
            ReduceMin(TBufTensorName(STORE_VALUE), TBufTensorName(STORE_VALUE), TBufTensorName(WORK_LOCAL), count + 1);
        }
        if (is_final)
        {
            QueAllocSimple(Y);
            DataCopy(LTensorName(Y), TBufTensorName(STORE_VALUE), 32 / sizeof(_DT_X));
            EnQue(Y);
        }
        QueFree(X);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount, bool is_final, int64_t big_index)
    {
        if (is_final)
        {
            DeQueLocal2Global(Y, big_index, 32 / sizeof(_DT_X));
        }
    }
};

#ifndef __CCE_KT_TEST__
extern "C" __aicore__ inline void run(GM_ADDR x, GM_ADDR y, int64_t dType, int64_t size, float p, float epsilon,
                                      CommonTiling &tiling, NdTensorSortByDimHelper *sortHelper = nullptr)
{
    if (p == 0)
    {
        AscendC::KernelP0<DTYPE_X, DTYPE_Y> op;
        op.Init(x, y, size, tiling, sortHelper);
        op.Process();
    }
    else if (p == 1)
    {
        AscendC::KernelP1<DTYPE_X, DTYPE_Y> op;
        op.Init(x, y, size, tiling, sortHelper);
        op.Process();
    }
    else if (p == 2)
    {
        AscendC::KernelP2<DTYPE_X, DTYPE_Y> op;
        op.Init(x, y, size, tiling, sortHelper);
        op.Process();
    }
    else if (p == float(1) / float(0))
    {
        AscendC::KernelPInf<DTYPE_X, DTYPE_Y> op;
        op.Init(x, y, size, tiling, sortHelper);
        op.Process();
    }
    else if (p == float(-1) / float(0))
    {
        AscendC::KernelPnInf<DTYPE_X, DTYPE_Y> op;
        op.Init(x, y, size, tiling, sortHelper);
        op.Process();
    }
    else
    {
        switch (dType)
        {
        case 0: {
            AscendC::KernelOtherFloat op;
            op.Init(x, y, size, p, epsilon, tiling, sortHelper);
            op.Process();
        }
        break;
        case 1: {
            AscendC::KernelOtherHalf op;
            op.Init(x, y, size, p, epsilon, tiling, sortHelper);
            op.Process();
        }
        break;
        }
    }
}
#endif
} // namespace AscendC
#ifndef __CCE_KT_TEST__
extern "C" __global__ __aicore__ void lp_norm_v2(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    CRATE_COMMON_TILING_SIMPLE;
    if (tiling_data.sortByDimDimSize != 0)
    {
        // 排序
        AscendC::NdTensorSortByDimHelper transposeHelper;
        transposeHelper.data_addr = x;
        transposeHelper.startIdx = commonTiling.startIdx;
        transposeHelper.length = commonTiling.blockLength;
        transposeHelper.sortByDimPreSize = tiling_data.sortByDimPreSize;
        transposeHelper.sortByDimDimSize = tiling_data.sortByDimDimSize;
        transposeHelper.sortByDimDimLength = tiling_data.sortByDimDimLength;
        transposeHelper.sortByDimDimBlockLength = tiling_data.sortByDimDimBlockLength;
        AscendC::run(x, y, tiling_data.dType, tiling_data.sortByDimDimSize, tiling_data.p, tiling_data.epsilon,
                     commonTiling, &transposeHelper);
    }
    else
    {
        AscendC::run(x, y, tiling_data.dType, tiling_data.size, tiling_data.p, tiling_data.epsilon, commonTiling);
    }
}
#else
extern "C" __global__ __aicore__ void lp_norm_v2(GM_ADDR x, GM_ADDR y)
{
    AscendC::CommonTiling commonTiling;
    commonTiling.Init(0, 0, 0, 0, 1, 4, 16, 16);
    // 排序
    AscendC::NdTensorSortByDimHelper transposeHelper;
    transposeHelper.data_addr = x;
    transposeHelper.startIdx = 0;
    transposeHelper.length = 4;
    transposeHelper.sortByDimPreSize = 2;
    transposeHelper.sortByDimDimSize = 2;
    transposeHelper.sortByDimDimLength = 2;
    transposeHelper.sortByDimDimBlockLength = 2;
    AscendC::KernelP1<half, half> op;
    op.Init(x, y, 8, commonTiling, &transposeHelper);
    op.Process();
}
#endif